"""
旁路Low-Rank Adaptation (LoRA)适配器实现。

本模块提供了一个适配器类，用于在Transformer块的intermediate和output层之间
创建LoRA旁路连接，实现高效的参数学习。
"""

from typing import List, Optional, Dict, Any, Union
import torch
import torch.nn as nn

from continuallearning.interfaces.models.adapter import AdapterType
from continuallearning.registry import HOOK_ADAPTER_REGISTRY
from continuallearning.models.pefts.hooks._base_hook import BaseHookAdapter
from continuallearning.models.backbones.base import BaseBackbone
from continuallearning.models.pefts.common.config import BypassLoRAConfig
from continuallearning.models.pefts.lora.hooks.bypass_lora_hooks import (
    BypassLoRAHookManager,
    BypassLoRAHookFactory,
)

# 替换标准logging库为自定义日志系统
from continuallearning.utils.logging import get_logger

# 使用自定义日志系统创建记录器
logger = get_logger(__name__)


@HOOK_ADAPTER_REGISTRY.register("bypass_lora")
class BypassLoRAAdapter(BaseHookAdapter):
    """
    在Intermediate和Output层之间实现低秩适配(LoRA)旁路连接的适配器。

    这个适配器允许你创建一个旁路连接，接收intermediate层的输入，并应用LoRA变换后
    将结果添加到output层的输出，实现更高效的参数学习。

    Args:
        backbone: 要适配的backbone模型
        rank: 低秩矩阵的维度
        alpha: 缩放因子
        dropout: Dropout概率
        intermediate_modules: 标识intermediate层的模块名模式列表
        output_modules: 标识output层的模块名模式列表
        task_specific: 是否为每个任务使用单独的权重
        num_tasks: 任务数量(当task_specific=True时需要)
        adapter_type: 适配器类型(通常为PRE_BACKBONE)
    """

    def __init__(
        self,
        backbone: BaseBackbone,
        rank: int = 4,
        alpha: float = 8.0,
        dropout: float = 0.1,
        intermediate_modules: Optional[List[str]] = None,
        output_modules: Optional[List[str]] = None,
        task_specific: bool = False,
        num_tasks: int = 10,
        adapter_type: AdapterType = AdapterType.PRE_BACKBONE,
        config: Optional[Union[Dict[str, Any], BypassLoRAConfig]] = None,
    ):
        # 初始化基类
        super().__init__(
            backbone=backbone,
            adapter_type=adapter_type,
            task_specific=task_specific,
            num_tasks=num_tasks,
        )

        # 处理配置：优先使用显式传递的配置对象，否则根据参数创建
        if config is not None:
            if isinstance(config, dict):
                self.config = BypassLoRAConfig(**config)
            else:
                self.config = config
        else:
            # 创建配置对象，保持向后兼容性
            self.config = BypassLoRAConfig(
                rank=rank,
                alpha=alpha,
                dropout=dropout,
                intermediate_modules=intermediate_modules
                or ["*.intermediate.dense", "*intermediate*"],
                output_modules=output_modules or ["*.output.dense", "*output*"],
            )

    def _create_hook_manager(
        self, model: nn.Module, task_id: int
    ) -> BypassLoRAHookManager:
        """
        为特定任务创建旁路LoRA Hook管理器。

        Args:
            model: 要应用hooks的模型
            task_id: 任务标识符

        Returns:
            BypassLoRAHookManager: 用于该任务的hook管理器
        """
        # 创建任务特定配置，保留所有模式参数
        task_config = BypassLoRAConfig(
            rank=self.config.rank,
            alpha=self.config.alpha,
            dropout=self.config.dropout,
            intermediate_modules=self.config.intermediate_modules,
            output_modules=self.config.output_modules,
            task_id=task_id,
        )

        # 更新hook工厂的配置，从配置对象获取模式参数
        task_hook_factory = BypassLoRAHookFactory(
            config=task_config,
        )

        # 创建并返回旁路hook管理器，从配置对象获取模式参数
        return BypassLoRAHookManager(
            model=model,
            intermediate_patterns=task_config.intermediate_modules,
            output_patterns=task_config.output_modules,
            hook_factory=task_hook_factory,
            config=task_config,
        )

    @property
    def lora_config(self) -> BypassLoRAConfig:
        """
        提供向后兼容的访问器，保持与现有代码的兼容性

        Returns:
            BypassLoRAConfig: 当前配置
        """
        return self.config

    def adapt_representation(
        self,
        hidden_states: torch.Tensor,
        task_id: Optional[int] = None,
        layer_idx: Optional[int] = None,
        **kwargs,
    ) -> torch.Tensor:
        """
        应用适配到模型表示(隐藏状态)。

        对于旁路LoRA适配器，与普通LoRA类似，适配发生在权重层级而非直接操作隐藏状态。

        Args:
            hidden_states: 输入隐藏状态(不变)
            task_id: 可选任务标识符
            layer_idx: 可选层索引
            **kwargs: 额外参数

        Returns:
            torch.Tensor: 原始输入隐藏状态(不变)
        """
        # 如果指定了task_id，确保激活正确的任务
        if task_id is not None and task_id != self.current_task:
            self.prepare_task(task_id)

        # 对于LoRA，适配发生在权重层级，通过hooks，不直接修改隐藏状态
        return hidden_states
